from time import time
import json
from channels.generic.websocket import AsyncWebsocketConsumer
from django_redis import get_redis_connection
from django.db import close_old_connections

def redis_cli():
    return get_redis_connection('default')

# 一个路径的URL 对应一个Consumer(或者一个用户也对应一个Consumer)
class SymbolSubscribe(AsyncWebsocketConsumer):
    """订阅器，常用于订阅某一类信息。当有新消息时，推送至客户端。
    客户端可能需要向服务端发送心跳包，维持在线状态。
    订阅者Redis Key：subscribers:stream_name:symbol 类型为SortedSet
    命令队列：commands: 类型为列表。
    """
    stream_name = 'bitmex_symbol'  # 监听的消息类型，结合参数如 ETC@trade
    heartbeat_enable = False  # 是否开启心跳
    user_sub_symbols = {}

    @property
    def extra_data(self):
        return

    def subscribe_message(self, symbol):
        """订阅推送消息
        """
        cli = redis_cli()
        cli.rpush(
            f'event_bus:bitmex:{self.stream_name}',
            json.dumps({
                'action': 'subscribe',
                'client': self.channel_name,
                'symbol': symbol,
                'time': time(),
                'stream': self.stream_name,
            })
        )

    def unsubscribe_message(self, symbol):
        """退订推送消息
        """
        cli = redis_cli()
        cli.rpush(
            f'event_bus:bitmex:{self.stream_name}',
            json.dumps({
                'action': 'unsubscribe',
                'stream': self.stream_name,
                'client': self.channel_name,
                'symbol': symbol,
                'time': time(),
            })
        )

    async def connect(self):
        """先鉴权，再订阅
        """
        self.user = self.scope['user']
        print(self.user)
        if not self.user.is_authenticated:
            await self.close()
        else:
            close_old_connections()  # 先关掉旧连接，以防出现连接过期异常
            if not self.user.account:
                # 如果不属于一个关易所帐号，关闭
                await self.close()
                return

        await self.accept()
        self.update_active_time()  # 更新活跃时间。

    def update_active_time(self):
        """更新自己最后活跃时间
        """
        cli = redis_cli()
        set_key = 'subscribers:%s:%s' % (
            self.stream_name,
            # symbol,
            # self.user.account.platform.slug,
            'bitmex'
        )
        cli.zadd(set_key, time(), self.channel_name)

    async def disconnect(self, close_code):
        """断开连接，退订队例
        """
        cli = redis_cli()
        set_key = 'subscribers:%s:%s' % (
            self.stream_name,
            'bitmex'
        )
        cli.zrem(set_key, self.channel_name)
        # if self.scope["user"] in self.user_sub_symbols:
        #     self.user_sub_symbols.pop(self.scope["user"])

    async def receive(self, text_data):
        """处理客户端发过来的消息，用于订阅bitmex合约，或者刷新时间戳
        1. {"action": "subscribe_symbol", "data" : symbol}
        2. {"action": "ping", "data": ""}

        """
        # self.update_active_time()
        # self.subscribe_message()
        user = self.scope['user']
        try:
            data = json.loads(text_data)
        except Exception:
            pass
        else:
            if isinstance(data, dict):
                if data.get('action') == "subscribe_symbol":
                    # ToDo 客户端需确认币对已订阅
                    symbol = data["data"].upper()
                    if user not in self.user_sub_symbols:
                        self.user_sub_symbols[user] = set(symbol)
                    else:
                        self.user_sub_symbols[user].add(symbol)
                    self.subscribe_message(symbol)
                elif data.get('action') == "subscribe_symbol":
                    symbol = data["data"].upper()
                    if user in self.user_sub_symbols:
                        self.user_sub_symbols[user].remove(symbol)
                    self.unsubscribe_message(symbol)

        self.update_active_time()
        await self.send(text_data=text_data)

    async def update_message(self, event):
        """从队列得到消息，在此方法中，需要把消息发送到客户端。
        """
        message = event['message']
        # 筛选出客户端订阅的币对信息后推送
        print(message)
        if self.user_sub_symbols.get(self.scope["user"]):
            for symbol in self.user_sub_symbols[self.scope["user"]]:
                if symbol in message:
                    await self.send(text_data=message)
                    break


class AccountSubscribe(AsyncWebsocketConsumer):
    """监听个人帐号更新
    """
    stream_name = 'bitmex_user'
    user_sub_topics = set()

    def subscribe_topic(self, topic):
        """订阅账户topic
        """
        cli = redis_cli()
        symbol = self.scope["user"].account.id
        cli.rpush(
            f'bitmex:user:{symbol}',
            json.dumps({
                'action': 'subscribe',
                'client': self.channel_name,
                'topic': topic,
                'time': time(),
                'stream': self.stream_name,
            })
        )

    def update_active_time(self):
        """更新自己最后活跃时间
        """
        cli = redis_cli()
        symbol = self.scope["user"].account.id
        set_key = 'subscribers:%s:%s' % (
            symbol,
            'bitmex'
        )
        print(set_key)
        # 移除掉之前的订阅的键值，多个键会严重影响redis的推送
        cli.zremrangebyrank(set_key, 0, -1)
        cli.zadd(set_key, time(), self.channel_name)

    async def connect(self):
        """先鉴权，再订阅
        """
        self.user = self.scope['user']
        print(self.user.name)
        print(self.user.account.id)
        if not self.user.is_authenticated:
            await self.close()
        else:
            close_old_connections()  # 先关掉旧连接，以防出现连接过期异常
            if not self.user.account:
                # 如果不属于一个关易所帐号，关闭
                await self.close()
                return
            await self.accept()
            # self.subscribe_message()
            self.update_active_time()  # 更新活跃时间。

    async def update_message(self, event):
        """从队列得到消息，在此方法中，需要把消息发送到客户端。
        """
        message = event['message']
        # 筛选出客户端订阅的币对信息后推送
        print(message)
        await self.send(text_data=message)

    async def disconnect(self, close_code):
        """断开连接，退订队例
        """
        cli = redis_cli()
        symbol = self.scope["user"].account.id
        set_key = 'subscribers:%s:%s' % (
            symbol,
            'bitmex'
        )
        cli.zrem(set_key, self.channel_name)

    async def receive(self, text_data):
        """处理客户端发过来的消息，用于订阅其他账户topic，或者刷新时间戳
        1. {"action": "subscribe_topic", "data" : topic}
        2. {"action": "ping", "data": ""}

        """
        user = self.scope['user']
        try:
            data = json.loads(text_data)
        except Exception:
            pass
        else:
            if isinstance(data, dict) and data.get('action') == "subscribe_topic":
                # ToDo 客户端需确认币对已订阅
                topic = data["data"].upper()
                if user not in self.user_sub_topics:
                    self.user_sub_topics.add(topic)
                else:
                    self.user_sub_topics.remove(topic)
                self.subscribe_topic(topic)
        self.update_active_time()
        await self.send(text_data=text_data)
